import sys, time, os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

from PyQt5.QtWidgets import QWidget, QApplication, QFileDialog, QMessageBox
from pyqt5_plugins.examplebuttonplugin import QtGui
import numpy as np
import uiTools.saveAsExl, uiTools.feature, uiTools.show
import aug_tools.augmentations, aug_tools.transforms
from tqdm import tqdm
import cv2
from PIL import Image
import detect
import shutil
import skimage
from torchvision import transforms
from PIL import Image
from typing import List, Optional, Tuple, Union
import gen_images, joint

# 全局窗口app
app = QApplication(sys.argv)


class MyWindows(new.Ui_widget, QWidget):
    def __init__(self):
        super(MyWindows, self).__init__()
        # 初始化UI
        self.setupUi(self)
        # 初始绑定各个按钮的功能
        self.handle_buttons()
        # tips信息
        # self.selectedFile.setText('暂无,请点击左侧按钮选择路径')
        # self.selectedSavePath.setText('暂无,请点击左侧按钮选择路径')
        # 检测类型 0为默认  单张检测为1  文件夹多张检测为2
        self.detectType = 0
        # 当选择检测一个文件夹的时候，其中保存当前文件夹中所有以bmp,jpg,png结尾的文件名称
        self.picArray = []

    # 初始绑定各个按钮的功能
    def handle_buttons(self):
        # 选择待检测图片
        self.selectFile.clicked.connect(self.onSelectPicFile)
        # 选择待检测图片文件
        self.selectFileDir.clicked.connect(self.onSelectPicFileDir)
        # 选择结果保存路径
        self.selectSavePath.clicked.connect(self.onSelectSavePath)
        # 开始检测
        self.detect.clicked.connect(self.onDetectPic)

        # 选择对比度图片文件夹
        self.pushButtonContrast.clicked.connect(self.onSelectContrastFileDir)
        # 选择对比度增强保存路径
        self.pushButtonContrastSave.clicked.connect(self.onSelectContrastSavePath)

        self.pushButtonContrastLabel.clicked.connect(self.onSelectContrastLabelPath)
        # 开始对比度增强
        self.contrastConfirm.clicked.connect(self.onContrastPic)

        # 选择翻转图片文件夹
        self.pushButtonReverse.clicked.connect(self.onSelectReverseFileDir)
        # 选择翻转增强保存路径
        self.pushButtonReverseSave.clicked.connect(self.onSelectReverseSavePath)
        # 选择翻转增强标签路径
        self.pushButtonReverseLabel.clicked.connect(self.onSelectReverseLabelPath)
        # 开始旋转增强
        self.reverseConfirm.clicked.connect(self.onReversePic)

        # 选择旋转图片文件夹
        self.pushButtonVein.clicked.connect(self.onSelectVeinFileDir)
        # 选择旋转增强保存路径
        self.pushButtonVeinSave.clicked.connect(self.onSelectVeinSavePath)
        # 选择旋转增强标签路径
        self.pushButtonVeinLabel.clicked.connect(self.onSelectVeinLabelPath)

        self.veinConfirm.clicked.connect(self.onVeinPic)

        self.pushButtonGAN.clicked.connect(self.onSelectGANSavePath)

        self.ganConfirm.clicked.connect(self.onGANPic)

    # 选择待检测图片
    def onSelectPicFile(self):
        # 选择待检测的图片,展示在待检测区
        filename, filetype = QFileDialog.getOpenFileName(self, '请选择待检图片')
        self.singlePicPath = filename
        # 显示选择路径的指定图片
        jpg = QtGui.QPixmap(filename).scaled(self.imageShow.width(), self.imageShow.height())
        self.selectedFile.setText(filename)
        # self.selectedFile.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        self.imageShow.setPixmap(jpg)
        # 调整检测类型为单张
        self.detectType = 1

    # 选择待检测图片文件
    def onSelectPicFileDir(self):
        # 选择文件夹路径
        self.selectedFile.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        # 调整检测类型为文件夹多张
        self.detectType = 2
        # 选择文件夹时 仅展示文件夹中第一张图片
        self.picArray = []

        # 判断当前文件中是否存在以jpg、bmp、png为结尾的文件
        for root, dirs, files in os.walk(self.selectedFile.text()):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                # 用文件后缀名判断当前文件夹中是否有图片类型
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)

        # 当文件夹中图片的个数大于0的时候才会展示一张出来
        if len(self.picArray) > 0:
            self.DirFirstPicPath = self.selectedFile.text() + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.imageShow.width(), self.imageShow.height())
            self.imageShow.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)
            self.selectedFile.setText('暂无,请点击左侧按钮选择路径')

    # 选择结果保存路径 是个文件夹就行
    def onSelectSavePath(self):
        self.selectedSavePath.setText(QFileDialog.getExistingDirectory())

    # 开始检测
    def onDetectPic(self):
        # 判断图片和保存路径是否选择 没选择就显示提示信息并不往下执行
        if self.selectedSavePath.text() == '暂无,请点击左侧按钮选择路径' and self.detectType == 1:
            QMessageBox.information(self, '提示', '请选择待检测图片路径!', QMessageBox.Yes)

        if self.selectedSavePath.text() == '暂无,请点击左侧按钮选择路径' and self.detectType == 2:
            QMessageBox.information(self, '提示', '请选择结果保存路径!', QMessageBox.Yes)

        # detectType为1时候，进行单张图片的检测
        elif self.selectedSavePath.text() != '暂无,请点击左侧按钮选择路径' and self.detectType == 1:
            # 图片保存文件夹设置为当前时间
            t = time.localtime()
            fileDirName = time.strftime("%Y-%m-%d-%H-%M-%S", t)
            # 在yolo的detect文件中重新创建了一个方法outSide,传入指定参数同样可以进行检测
            # 参数 source=待检测图片 savePath=结果保存路径 fileDirName=结果要保存在一个文件夹中
            detect.outSide(
                source=self.selectedFile.text(),
                savePath=self.selectedSavePath.text(),
                name=str(fileDirName)
            )

            # 通过检测结果保存的路径,在待检测区展示检测结果图
            picName = self.singlePicPath.split("/")
            picNameLast = picName[-1]
            showImagePath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + picNameLast
            jpg = QtGui.QPixmap(showImagePath).scaled(self.imageShowDetect.width(), self.imageShowDetect.height())
            self.imageShowDetect.setPixmap(jpg)

            # 统计检测出的目标个数
            countGranularSludge = 0
            countEpistylis = 0
            countFloc = 0
            countRotifer = 0
            countFilamentousbacteria = 0

            # # 一张图中所有颗粒污泥的实际直径
            # averageAd = []
            # averageFd = []
            # 循环读取裁剪出来的目标个数,统计检测结果
            for root, dirs, files in os.walk(
                    self.selectedSavePath.text() + "/" + fileDirName + "./crops/Granular Sludge/"):
                countGranularSludge = len(files)
                # for file in files:
                #     # 此处进行单张图片的特征计算
                #
                #     # 分割后的绝对路径
                #     absPath = self.selectedSavePath.text() + "/" + fileDirName + "./crops/Granular Sludge/" + file
                #     # 读取图片特征计算
                #
                #     # 计算每张图的最大颗粒粒径
                #     averageAd.append(uiTools.feature.featureCalculateAd(absPath))
                #     # 计算颗粒污泥分形
                #     averageFd.append(uiTools.feature.featureCalculateFd(absPath))

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Epistylis/"):
                countEpistylis = len(files)

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Rotifer/"):
                countRotifer = len(files)

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Floc/"):
                countFloc = len(files)

            for root, dirs, files in os.walk(
                    self.selectedSavePath.text() + "/" + fileDirName + "./crops/Filamentous bacteria/"):
                countFilamentousbacteria = len(files)

            # 拼接保存路径
            excelPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName)

            # 把目标个数统计结果保存至excel
            uiTools.saveAsExl.saveDetectAsExl(excelPath,
                                              countGranularSludge,
                                              countEpistylis,
                                              countFloc,
                                              countRotifer,
                                              countFilamentousbacteria,
                                              )

            # 把颗粒污泥等效直径保存至excel
            # uiTools.saveAsExl.saveRealAdAsExl(excelPath, averageAd, averageFd)

            # 计算所以颗粒污泥的平均粒径
            # totleavead = 0
            # for ele in range(0, len(averageAd)):
            #     totleavead = totleavead + averageAd[ele]
            #
            # totleavefd = 0
            # for ele in range(0, len(averageFd)):
            #     totleavefd = totleavefd + averageFd[ele]
            # 展示检测结果
            # self.statisticTitle.setText('检测结果统计如下:')
            # self.statistic.setText('颗粒污泥个数:' + str(countGranularSludge) + '个\n' +
            #                        '累枝虫个数:' + str(countEpistylis) + '个\n' +
            #                        '轮虫个数:' + str(countRotifer) + '个\n' +
            #                        '絮体个数:' + str(countFloc) + '个\n' +
            #                        '丝状菌个数:' + str(countFilamentousbacteria) + '个\n\n' +
            #                        '颗粒污泥平均粒径:' + str(round(totleavead / countGranularSludge, 2)) + ' um\n' +
            #                        '颗粒污泥平均分形:' + str(round(totleavefd / countGranularSludge, 4)) + '\n' +
            #                        '颗粒污泥粒径标准差:' + str(round(np.std(averageAd), 3)) + ' \n' +
            #                        '颗粒污泥分形标准差:' + str(round(np.std(averageFd), 3)) + ' \n')

            # self.statistic.setText('颗粒污泥个数:' + str(countGranularSludge) + '个\n' +
            #                        '累枝虫个数:' + str(countEpistylis) + '个\n' +
            #                        '轮虫个数:' + str(countRotifer) + '个\n' +
            #                        '絮体个数:' + str(countFloc) + '个\n' +
            #                        '丝状菌个数:' + str(countFilamentousbacteria) + '个\n\n')

            # 图表保存路径
            # figurePath = excelPath + './' + fileDirName
            # uiTools.show.savePic(averageAd, figurePath)
            # # 粒径分布图展示
            # figPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + fileDirName + '.png'
            # StatisticJpg = QtGui.QPixmap(figPath).scaled(self.imageShowStatis.width(), self.imageShowStatis.height())
            # self.imageShowStatis.setPixmap(StatisticJpg)

            figurePath = excelPath + './' + fileDirName
            uiTools.show.savePic(excelPath, figurePath)

            figPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + fileDirName + '.png'
            StatisticJpg = QtGui.QPixmap(figPath).scaled(self.imageShowStatis.width(), self.imageShowStatis.height())
            self.imageShowStatis.setPixmap(StatisticJpg)

        # detectType为2时候，进行一个文件夹中所有图片的检测
        elif self.selectedSavePath.text() != '暂无,请点击左侧按钮选择路径' and self.detectType == 2:
            # 图片保存文件夹设置为当前时间
            t = time.localtime()
            fileDirName = time.strftime("%Y-%m-%d-%H-%M-%S", t)
            # 在yolo的detect文件中重新创建了一个方法outSide,传入指定参数同样可以进行检测
            # 参数 source=待检测图片 savePath=结果保存路径 fileDirName=结果要保存在一个文件夹中
            detect.outSide(
                source=self.selectedFile.text(),
                savePath=self.selectedSavePath.text(),
                name=str(fileDirName)
            )

            # 只展示出待检测文件中第一张图片的检测结果，展示在界面上
            showImagePath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + self.picArray[0]
            jpg = QtGui.QPixmap(showImagePath).scaled(self.imageShowDetect.width(), self.imageShowDetect.height())
            self.imageShowDetect.setPixmap(jpg)

            # 统计文件夹中所有图片检测出的目标个数
            countGranularSludge = 0
            countEpistylis = 0
            countFloc = 0
            countRotifer = 0
            countFilamentousbacteria = 0

            # averageAd = []
            # averageFd = []
            detectedPicCount = 0
            for root, dirs, files in os.walk(self.selectedFile.text()):
                detectedPicCount = len(files)
                # print('一共有',detectedPicCount,'张图片被检测')
            for root, dirs, files in os.walk(
                    self.selectedSavePath.text() + "/" + fileDirName + "./crops/Granular Sludge/"):
                countGranularSludge = len(files)
                # for file in files:
                #     # 此处进行单张图片的特征计算
                #
                #     # 分割后的绝对路径
                #     absPath = self.selectedSavePath.text() + "/" + fileDirName + "./crops/Granular Sludge/" + file
                #     # 读取图片特征计算
                #
                #     # 计算每张图的最大颗粒粒径
                #     averageAd.append(uiTools.feature.featureCalculateAd(absPath))
                #     # 计算颗粒污泥分形
                #     averageFd.append(uiTools.feature.featureCalculateFd(absPath))

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Epistylis/"):
                countEpistylis = len(files)

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Rotifer/"):
                countRotifer = len(files)

            for root, dirs, files in os.walk(self.selectedSavePath.text() + "/" + fileDirName + "./crops/Floc/"):
                countFloc = len(files)

            for root, dirs, files in os.walk(
                    self.selectedSavePath.text() + "/" + fileDirName + "./crops/Filamentous bacteria/"):
                countFilamentousbacteria = len(files)

            # 拼接保存路径
            excelPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName)

            # 把目标个数统计结果保存至excel
            uiTools.saveAsExl.saveDetectAsExl(excelPath,
                                              countGranularSludge,
                                              countEpistylis,
                                              countFloc,
                                              countRotifer,
                                              countFilamentousbacteria,
                                              )

            # 把颗粒污泥等效直径保存至excel
            # uiTools.saveAsExl.saveRealAdAsExl(excelPath, averageAd, averageFd)
            #
            # # 计算所以颗粒污泥的平均粒径
            # totleavead = 0
            # for ele in range(0, len(averageAd)):
            #     totleavead = totleavead + averageAd[ele]
            #
            # totleavefd = 0
            # for ele in range(0, len(averageFd)):
            #     totleavefd = totleavefd + averageFd[ele]
            # 展示检测结果
            # self.statisticTitle.setText('检测结果统计如下:')
            # self.statistic.setText('全部颗粒污泥个数:' + str(countGranularSludge) + '个\n' +
            #                        '全部累枝虫个数:' + str(countEpistylis) + '个\n' +
            #                        '全部轮虫个数:' + str(countRotifer) + '个\n' +
            #                        '全部絮体个数:' + str(countFloc) + '个\n' +
            #                        '全部丝状菌个数:' + str(countFilamentousbacteria) + '个\n\n' +
            #                        '平均颗粒污泥个数:' + str(countGranularSludge/detectedPicCount) + '个\n' +
            #                        '平均累枝虫个数:' + str(countEpistylis/detectedPicCount) + '个\n' +
            #                        '平均轮虫个数:' + str(countRotifer/detectedPicCount) + '个\n' +
            #                        '平均絮体个数:' + str(countFloc/detectedPicCount) + '个\n' +
            #                        '平均丝状菌个数:' + str(countFilamentousbacteria/detectedPicCount) + '个\n\n')

            # self.statistic.setText('全部颗粒污泥个数:' + str(countGranularSludge) + '个\n' +
            #                        '全部累枝虫个数:' + str(countEpistylis) + '个\n' +
            #                        '全部轮虫个数:' + str(countRotifer) + '个\n' +
            #                        '全部絮体个数:' + str(countFloc) + '个\n' +
            #                        '全部丝状菌个数:' + str(countFilamentousbacteria) + '个\n\n' +
            #                        '平均颗粒污泥个数:' + str(countGranularSludge / detectedPicCount) + '个\n' +
            #                        '平均累枝虫个数:' + str(countEpistylis / detectedPicCount) + '个\n' +
            #                        '平均轮虫个数:' + str(countRotifer / detectedPicCount) + '个\n' +
            #                        '平均絮体个数:' + str(countFloc / detectedPicCount) + '个\n' +
            #                        '平均丝状菌个数:' + str(countFilamentousbacteria / detectedPicCount) + '个\n\n' +
            #                        '颗粒污泥平均粒径:' + str(round(totleavead / countGranularSludge, 2)) + ' um\n' +
            #                        '颗粒污泥平均分形:' + str(round(totleavefd / countGranularSludge, 4)) + '\n' +
            #                        '颗粒污泥粒径标准差:' + str(round(np.std(averageAd), 3)) + ' \n' +
            #                        '颗粒污泥分形标准差:' + str(round(np.std(averageFd), 3)) + ' \n')

            # 图表保存路径
            # figurePath = excelPath + './' + fileDirName
            # uiTools.show.savePic(averageAd, figurePath)
            # 粒径分布图展示
            # figPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + fileDirName + '.png'
            # StatisticJpg = QtGui.QPixmap(figPath).scaled(self.imageShowStatis.width(), self.imageShowStatis.height())
            # self.imageShowStatis.setPixmap(StatisticJpg)
            figurePath = excelPath + './' + fileDirName
            epath = 'EXCEL/results.xls'
            empath = 'EXCEL/results_mean.xls'
            uiTools.show.savePicMulti(epath, empath, figurePath)

            figPath = str(self.selectedSavePath.text()) + '/' + str(fileDirName) + '/' + fileDirName + '.png'
            StatisticJpg = QtGui.QPixmap(figPath).scaled(self.imageShowStatis.width(), self.imageShowStatis.height())
            self.imageShowStatis.setPixmap(StatisticJpg)

    def onSelectContrastFileDir(self):
        # 选择文件夹路径
        self.contrastPath.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        # 调整检测类型为文件夹多张
        self.detectType = 2
        # 选择文件夹时 仅展示文件夹中第一张图片
        self.picArray = []

        # 判断当前文件中是否存在以jpg、bmp、png为结尾的文件
        for root, dirs, files in os.walk(self.contrastPath.text()):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                # print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)
                    print(self.picArray)

        # 当文件夹中图片的个数大于0的时候才会展示一张出来
        if len(self.picArray) > 0:
            self.DirFirstPicPath = self.contrastPath.text() + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_pre_contrast.width(),
                                                             self.aug_pre_contrast.height())
            self.aug_pre_contrast.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    def onSelectContrastLabelPath(self):
        # 选择数据集标签所在文件夹
        self.contrastPathLabel.setText(QFileDialog.getExistingDirectory())

    def onSelectContrastSavePath(self):
        # 选择结果保存路径
        self.contrastPathSave.setText(QFileDialog.getExistingDirectory())

    # 对比度增广
    def onContrastPic(self):
        # 获取文件下属性为imgProperty的所有文件
        def GetImgNameByEveryDir(file_dir, imgProperty):
            FileName = []
            for root, dirs, files in os.walk(file_dir):
                for file in files:
                    if os.path.splitext(file)[1] in imgProperty:
                        FileName.append(file)  # 保存图片名称
            return FileName

        def readBoxes(txt_path):
            boxes = []
            with open(txt_path) as file:
                txt_lines = file.readlines()
                for txt_line in txt_lines:
                    box = txt_line.rstrip().split(" ")
                    boxes.append([int(box[0]), float(box[1]), float(box[2]), float(box[3]), float(box[4])])

            return boxes

        image_path = self.contrastPath.text()
        label_path = self.contrastPathLabel.text()
        output_path = self.contrastPathSave.text()

        pho = str(self.photoPara.currentText())

        def get_class_value(className):
            if className == "对比度变换":
                return 0
            elif className == "亮度变换":
                return 1
            elif className == "饱和度变换":
                return 2

        phoId = get_class_value(pho)

        img_list = GetImgNameByEveryDir(image_path, ['.jpg', '.jpeg', '.bmp'])
        for img_name in tqdm(img_list):
            img_is_ok = 1
            boxes = []
            img_path = image_path + '\\' + img_name
            try:
                img = np.array(Image.open(img_path).convert('RGB'), dtype=np.uint8)
                img1 = cv2.imread(img_path)
            except Exception as e:
                print(f"could not read image '{img_path}'. ")
                img_is_ok = 0
            if img_is_ok:  # 如果图像存在，读取对应的标签文件
                label_name = img_name[:-3] + 'txt'
                txt_path = label_path + '\\' + label_name
                boxes = readBoxes(txt_path)
                print(boxes)

            if phoId == 0:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Contrast')):
                    os.makedirs(os.path.join(output_path, 'Contrast'))
                if not os.path.exists(os.path.join(output_path, 'Contrast', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Contrast', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Contrast', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Contrast', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Contrast', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Contrast
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Contrast_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Contrast', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Contrast', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                contrastPath = self.contrastPathSave.text() + '/' + 'Contrast' + '/' + 'aug_img'

            if phoId == 1:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Brightness')):
                    os.makedirs(os.path.join(output_path, 'Brightness'))
                if not os.path.exists(os.path.join(output_path, 'Brightness', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Brightness', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Brightness', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Brightness', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Brightness', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Brightness
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Brightness_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Brightness', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Brightness', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                contrastPath = self.contrastPathSave.text() + '/' + 'Brightness' + '/' + 'aug_img'

            elif phoId == 2:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Saturation')):
                    os.makedirs(os.path.join(output_path, 'Saturation'))
                if not os.path.exists(os.path.join(output_path, 'Saturation', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Saturation', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Saturation', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Saturation', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Saturation', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_HueAndSaturation
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Saturation_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Saturation', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Saturation', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                contrastPath = self.contrastPathSave.text() + '/' + 'Saturation' + '/' + 'aug_img'

        self.picArray = []
        # reversePath = self.reversePathSave.text() + '/' + 'reverse' + '/' + 'aug_img'
        for root, dirs, files in os.walk(contrastPath):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)

        if len(self.picArray) > 0:
            self.DirFirstPicPath = contrastPath + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_after_contrast.width(),
                                                             self.aug_after_contrast.height())
            self.aug_after_contrast.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    # 选择纹理增广数据集图片路径
    def onSelectVeinFileDir(self):
        # 选择文件夹路径
        self.veinPath.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        # 调整检测类型为文件夹多张
        self.detectType = 2
        # 选择文件夹时 仅展示文件夹中第一张图片
        self.picArray = []

        # 判断当前文件中是否存在以jpg、bmp、png为结尾的文件
        for root, dirs, files in os.walk(self.veinPath.text()):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)
                    # print(self.picArray)
        if len(self.picArray) > 0:
            self.DirFirstPicPath = self.veinPathSave.text() + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_after_vein.width(), self.aug_after_vein.height())
            self.aug_after_vein.setPixmap(jpg)
        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    # 选择纹理增广结果保存路径
    def onSelectVeinFileDir(self):
        # 选择文件夹路径
        self.veinPath.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        # 调整检测类型为文件夹多张
        self.detectType = 2
        # 选择文件夹时 仅展示文件夹中第一张图片
        self.picArray = []

        # 判断当前文件中是否存在以jpg、bmp、png为结尾的文件
        for root, dirs, files in os.walk(self.veinPath.text()):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)
                    # print(self.picArray)

        # 当文件夹中图片的个数大于0的时候才会展示一张出来
        if len(self.picArray) > 0:
            self.DirFirstPicPath = self.veinPath.text() + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_pre_vein.width(), self.aug_pre_vein.height())
            self.aug_pre_vein.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    def onSelectVeinLabelPath(self):
        self.veinPathLabel.setText(QFileDialog.getExistingDirectory())

    def onSelectVeinSavePath(self):
        self.veinPathSave.setText(QFileDialog.getExistingDirectory())

    # 纹理增广方法
    def onVeinPic(self):
        # 获取文件下属性为imgProperty的所有文件
        def GetImgNameByEveryDir(file_dir, imgProperty):
            FileName = []
            for root, dirs, files in os.walk(file_dir):
                for file in files:
                    if os.path.splitext(file)[1] in imgProperty:
                        FileName.append(file)  # 保存图片名称
            return FileName

        def readBoxes(txt_path):
            boxes = []
            with open(txt_path) as file:
                txt_lines = file.readlines()
                for txt_line in txt_lines:
                    box = txt_line.rstrip().split(" ")
                    boxes.append([int(box[0]), float(box[1]), float(box[2]), float(box[3]), float(box[4])])

            return boxes

        # 程序入口
        # --img_path为需要扩增的图像数据
        # note:图像和标签文件存在同一个文件夹
        # 标签坐标：中心点坐标和宽高（cX, cY, W, H）

        image_path = self.veinPath.text()
        label_path = self.veinPathLabel.text()
        output_path = self.veinPathSave.text()

        vein = str(self.VeinPara.currentText())

        def get_class_value(className):
            if className == "高斯噪声":
                return 0
            elif className == "高斯扰动":
                return 1

        veinId = get_class_value(vein)

        img_list = GetImgNameByEveryDir(image_path, ['.jpg', '.jpeg', '.bmp'])
        for img_name in tqdm(img_list):
            img_is_ok = 1
            boxes = []
            img_path = image_path + '\\' + img_name
            try:
                img = np.array(Image.open(img_path).convert('RGB'), dtype=np.uint8)
                img1 = cv2.imread(img_path)
            except Exception as e:
                print(f"could not read image '{img_path}'. ")
                img_is_ok = 0
            if img_is_ok:  # 如果图像存在，读取对应的标签文件
                label_name = img_name[:-3] + 'txt'
                txt_path = label_path + '\\' + label_name
                boxes = readBoxes(txt_path)
                print(boxes)

            if veinId == 0:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Noise')):
                    os.makedirs(os.path.join(output_path, 'Noise'))
                if not os.path.exists(os.path.join(output_path, 'Noise', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Noise', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Noise', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Noise', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Noise', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_GaussianNoise
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Noise_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Noise', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Noise', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                veinPath = self.veinPathSave.text() + '/' + 'Noise' + '/' + 'aug_img'


            elif veinId == 1:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Blur')):
                    os.makedirs(os.path.join(output_path, 'Blur'))
                if not os.path.exists(os.path.join(output_path, 'Blur', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Blur', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Blur', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Blur', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Blur', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Blur
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Blur_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Blur', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Blur', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                veinPath = self.veinPathSave.text() + '/' + 'Blur' + '/' + 'aug_img'

        self.picArray = []
        # reversePath = self.reversePathSave.text() + '/' + 'reverse' + '/' + 'aug_img'
        for root, dirs, files in os.walk(veinPath):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)

        if len(self.picArray) > 0:
            self.DirFirstPicPath = veinPath + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_after_vein.width(), self.aug_after_vein.height())
            self.aug_after_vein.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    # 选择图片翻转增广图片文件夹
    def onSelectReverseFileDir(self):
        # 选择文件夹路径
        self.reversePath.setText(QFileDialog.getExistingDirectory(self, '请选择待检文件夹'))
        # 调整检测类型为文件夹多张
        self.detectType = 2
        # 选择文件夹时 仅展示文件夹中第一张图片
        self.picArray = []

        # 判断当前文件中是否存在以jpg、bmp、png为结尾的文件
        for root, dirs, files in os.walk(self.reversePath.text()):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)
                    # print(self.picArray)

        # 当文件夹中图片的个数大于0的时候才会展示一张出来
        if len(self.picArray) > 0:
            self.DirFirstPicPath = self.reversePath.text() + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_pre_reverse.width(),
                                                             self.aug_pre_reverse.height())
            self.aug_pre_reverse.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    # 选择图片翻转增广标签文件夹
    def onSelectReverseLabelPath(self):
        self.reversePathLabel.setText(QFileDialog.getExistingDirectory())

    # 选择图片翻转增广结果保存路径
    def onSelectReverseSavePath(self):
        self.reversePathSave.setText(QFileDialog.getExistingDirectory())

    # 图片翻转方法
    def onReversePic(self):
        # 获取文件下属性为imgProperty的所有文件
        def GetImgNameByEveryDir(file_dir, imgProperty):
            FileName = []
            for root, dirs, files in os.walk(file_dir):
                for file in files:
                    if os.path.splitext(file)[1] in imgProperty:
                        FileName.append(file)  # 保存图片名称
            return FileName

        def readBoxes(txt_path):
            boxes = []
            with open(txt_path) as file:
                txt_lines = file.readlines()
                for txt_line in txt_lines:
                    box = txt_line.rstrip().split(" ")
                    boxes.append([int(box[0]), float(box[1]), float(box[2]), float(box[3]), float(box[4])])

            return boxes

        # 程序入口
        # --img_path为需要扩增的图像数据
        # note:图像和标签文件存在同一个文件夹
        # 标签坐标：中心点坐标和宽高（cX, cY, W, H）

        image_path = self.reversePath.text()
        label_path = self.reversePathLabel.text()
        output_path = self.reversePathSave.text()

        geo = str(self.FlipPara.currentText())

        def get_class_value(className):
            if className == "水平翻转":
                return 0
            elif className == "上下翻转":
                return 1
            elif className == "随机旋转":
                return 2

        geoId = get_class_value(geo)

        img_list = GetImgNameByEveryDir(image_path, ['.jpg', '.jpeg', '.bmp'])
        for img_name in tqdm(img_list):
            img_is_ok = 1
            boxes = []
            img_path = image_path + '\\' + img_name
            try:
                img = np.array(Image.open(img_path).convert('RGB'), dtype=np.uint8)
                img1 = cv2.imread(img_path)
            except Exception as e:
                print(f"could not read image '{img_path}'. ")
                img_is_ok = 0
            if img_is_ok:  # 如果图像存在，读取对应的标签文件
                label_name = img_name[:-3] + 'txt'
                txt_path = label_path + '\\' + label_name
                boxes = readBoxes(txt_path)
                print(boxes)

            if geoId == 0:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Fliplr')):
                    os.makedirs(os.path.join(output_path, 'Fliplr'))
                if not os.path.exists(os.path.join(output_path, 'Fliplr', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Fliplr', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Fliplr', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Fliplr', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Fliplr', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Fliplr
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Fliplr_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Fliplr', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Fliplr', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                reversePath = self.reversePathSave.text() + '/' + 'Fliplr' + '/' + 'aug_img'


            elif geoId == 1:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Flipud')):
                    os.makedirs(os.path.join(output_path, 'Flipud'))
                if not os.path.exists(os.path.join(output_path, 'Flipud', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Flipud', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Flipud', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Flipud', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Flipud', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Flipud
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Flipud_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Flipud', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Flipud', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                reversePath = self.reversePathSave.text() + '/' + 'Flipud' + '/' + 'aug_img'

            elif geoId == 2:

                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                if not os.path.exists(os.path.join(output_path, 'Rotate')):
                    os.makedirs(os.path.join(output_path, 'Rotate'))
                if not os.path.exists(os.path.join(output_path, 'Rotate', 'aug_img')):
                    os.makedirs(os.path.join(output_path, 'Rotate', 'aug_img'))
                if not os.path.exists(os.path.join(output_path, 'Rotate', 'aug_label')):
                    os.makedirs(os.path.join(output_path, 'Rotate', 'aug_label'))
                    # 在aug_label文件夹创建classes.txt文件
                    with open(os.path.join(output_path, 'Rotate', 'aug_label', 'classes.txt'), 'w') as f:
                        f.write('Granular sludge\n')
                        f.write('Epistylis\n')
                        f.write('Floc\n')
                        f.write('Rotifer\n')
                        f.write('Filamentous bacteria\n')

                transform = aug_tools.augmentations.AUGMENTATION_TRANSFORMS_Rotate
                boxes = np.array(boxes)
                temp_boxes = np.zeros_like(boxes)
                temp_boxes[:, :] = boxes[:, :]
                #
                # copy_num为对同一张图片扩充的张数
                copy_num = 1  # 增强张数
                for i in np.arange(copy_num):
                    new_img, bb_target = transform((img1, boxes))
                    save_name = img_name[:-4] + "_Rotate_" + str(i)
                    cv2.imwrite(os.path.join(output_path, 'Rotate', 'aug_img', save_name + '.bmp'), new_img)
                    txt_file = open(os.path.join(output_path, 'Rotate', 'aug_label', save_name + '.txt'), 'w')
                    for line in bb_target:
                        bb = str(int(line[0])) + ' ' + str(line[1]) + ' ' + str(line[2]) + ' ' + str(
                            line[3]) + ' ' + str(
                            line[4]) + '\n'
                        txt_file.write(bb)
                    txt_file.close()
                    boxes[:, :] = temp_boxes[:, :]
                reversePath = self.reversePathSave.text() + '/' + 'Rotate' + '/' + 'aug_img'

        self.picArray = []
        # reversePath = self.reversePathSave.text() + '/' + 'reverse' + '/' + 'aug_img'
        for root, dirs, files in os.walk(reversePath):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)

        if len(self.picArray) > 0:
            self.DirFirstPicPath = reversePath + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_after_reverse.width(),
                                                             self.aug_after_reverse.height())
            self.aug_after_reverse.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)

    # gan增广方法结果保存路径  gan增广不需要选择图片和标签  通过模型进行增广
    def onSelectGANSavePath(self):
        self.ganPathSave.setText(QFileDialog.getExistingDirectory())

    # gan增广方法
    def onGANPic(self):
        out = './out'
        out_path = self.ganPathSave.text()

        class_name = str(self.ganClassPara.currentText())
        network = './weights/' + class_name + '.pkl'

        if not os.path.exists(out_path):
            os.makedirs(out_path)
        if not os.path.exists(os.path.join(out_path, 'GAN_' + class_name)):
            os.makedirs(os.path.join(out_path, 'GAN_' + class_name))
        if not os.path.exists(os.path.join(out_path, 'GAN_' + class_name, 'aug_img')):
            os.makedirs(os.path.join(out_path, 'GAN_' + class_name, 'aug_img'))
        if not os.path.exists(os.path.join(out_path, 'GAN_' + class_name, 'aug_label')):
            os.makedirs(os.path.join(out_path, 'GAN_' + class_name, 'aug_label'))
            # 在aug_label文件夹创建classes.txt文件
            with open(os.path.join(out_path, 'GAN_' + class_name, 'aug_label', 'classes.txt'), 'w') as f:
                f.write('Granular sludge\n')
                f.write('Epistylis\n')
                f.write('Floc\n')
                f.write('Rotifer\n')
                f.write('Filamentous bacteria\n')

        out_image = out_path + '/GAN_' + class_name + '/aug_img'
        out_label = out_path + '/GAN_' + class_name + '/aug_label'

        def get_class_value(className):
            if className == "GranularSludge":
                return 0
            elif className == "Epistylis":
                return 1
            elif className == "Floc":
                return 2
            elif className == "Rotifer":
                return 3
            elif className == "FilamentousBacteria":
                return 4

        classId = get_class_value(class_name)

        BigNumber = int(self.ganNumPara.currentText())
        SmallNumber = BigNumber * 10
        arr = [i for i in range(1, SmallNumber + 1)]

        trunc = 1

        def clear_folder(folder_path):
            for filename in os.listdir(folder_path):
                file_path = os.path.join(folder_path, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                except Exception as e:
                    print('Failed to delete %s. Reason: %s' % (file_path, e))

        clear_folder(out)

        gen_images.generate_images(network_pkl=network,
                                   seeds=arr,
                                   truncation_psi=trunc,
                                   noise_mode='const',
                                   outdir=out,
                                   translate=(0, 0),
                                   rotate=0,
                                   class_idx=None)

        joint.create_large_image_and_yolo_labels(input_folder=out,
                                                 output_image_folder=out_image,
                                                 output_label_folder=out_label,
                                                 class_id=classId)

        self.picArray = []
        for root, dirs, files in os.walk(out_image):
            for file in files:
                # 检测的目标个数
                ext = os.path.splitext(file)[-1].lower()
                print(ext)
                if ext == '.jpg' or ext == '.bmp' or ext == '.png':
                    self.picArray.append(file)

        if len(self.picArray) > 0:
            self.DirFirstPicPath = out_image + '/' + self.picArray[0]
            # 显示文件夹中的第一个图片
            jpg = QtGui.QPixmap(self.DirFirstPicPath).scaled(self.aug_after_gan.width(), self.aug_after_gan.height())
            self.aug_after_gan.setPixmap(jpg)

        else:
            QMessageBox.information(self, '提示', '当前选择文件夹中不存在以jpg、bmp、png结尾的图片文件，请重新选择！', QMessageBox.Yes)


my_windows = MyWindows()  # 实例化对象
my_windows.show()  # 显示窗口

sys.exit(app.exec_())
